{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from math import log2\n",
    "\n",
    "from functools import partial\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "from torch import nn, einsum\n",
    "from einops import rearrange\n",
    "from kornia.filters import filter2d\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def exists(val):\n",
    "    return val is not None\n",
    "\n",
    "def image_noise(n, im_size, device):\n",
    "    return torch.FloatTensor(n, 1, im_size, im_size).uniform_(0., 1.).to(device)\n",
    "\n",
    "def leaky_relu(p=0.2):\n",
    "    return nn.LeakyReLU(p, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Flatten(nn.Module):\n",
    "    def forward(self, x):\n",
    "        return x.reshape(x.shape[0], -1)\n",
    "\n",
    "class Blur(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        f = torch.Tensor([1, 2, 1])\n",
    "        self.register_buffer('f', f)\n",
    "    def forward(self, x):\n",
    "        f = self.f\n",
    "        f = f[None, None, :] * f [None, :, None]\n",
    "        return filter2d(x, f, normalized=True)\n",
    "    \n",
    "    \n",
    "class ConstantInput(nn.Module):\n",
    "    def __init__(self, channel, size=4):\n",
    "        super().__init__()\n",
    "\n",
    "        self.input = nn.Parameter(torch.randn(1, channel, size, size))\n",
    "\n",
    "    def forward(self, input):\n",
    "        batch = input.shape[0]\n",
    "        out = self.input.repeat(batch, 1, 1, 1)\n",
    "\n",
    "        return out\n",
    "    \n",
    "\n",
    "class GeneratorBlock(nn.Module):\n",
    "    def __init__(self, input_channels, filters, upsample=True):\n",
    "        super().__init__()\n",
    "        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False) if upsample else None\n",
    "\n",
    "        self.conv1 = nn.Conv2d(input_channels, filters, 3, 1, 1)\n",
    "        self.conv2 = nn.Conv2d(filters, filters, 3, 1, 1)\n",
    "\n",
    "        self.activation = leaky_relu()\n",
    "\n",
    "    def forward(self, x):\n",
    "        if exists(self.upsample):\n",
    "            x = self.upsample(x)\n",
    "\n",
    "        x = self.conv1(x)\n",
    "        x = self.activation(x)\n",
    "\n",
    "        x = self.conv2(x)\n",
    "        x = self.activation(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class Generator(nn.Module):\n",
    "    def __init__(\n",
    "            self,\n",
    "            image_size,\n",
    "            network_capacity=16,\n",
    "            fmap_max = 512,\n",
    "            latent_dim=512,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.num_layers = int(log2(image_size) - 1)\n",
    "        self.latent_dim = latent_dim\n",
    "        \n",
    "        filters = [network_capacity * (2 ** (i + 1)) for i in range(self.num_layers)][::-1]\n",
    "        set_fmap_max = partial(min, fmap_max)\n",
    "        filters = list(map(set_fmap_max, filters))\n",
    "        init_channels = filters[0]\n",
    "        \n",
    "        filters = [init_channels, *filters]\n",
    "        in_out_pairs = zip(filters[:-1], filters[1:])\n",
    "\n",
    "        self.initial_block = nn.ConvTranspose2d(latent_dim, init_channels, 4, 1, 0, bias=False)\n",
    "        self.initial_conv = nn.Conv2d(filters[0], filters[0], 3, padding=1)\n",
    "        \n",
    "        self.conv1 = nn.Conv2d(1, init_channels, 3, padding=1)\n",
    "        \n",
    "        self.blocks = nn.ModuleList([])\n",
    "        \n",
    "        for ind, (in_chan, out_chan) in enumerate(in_out_pairs):\n",
    "            not_first = ind != 0\n",
    "            # not_last = ind != (self.num_layers - 1)\n",
    "            \n",
    "            block = GeneratorBlock(\n",
    "                in_chan,\n",
    "                out_chan,\n",
    "                upsample=not_first,\n",
    "            )\n",
    "            self.blocks.append(block)\n",
    "            \n",
    "        self.to_mri = nn.Conv2d(\n",
    "            in_channels=filters[-1], \n",
    "            out_channels=1, \n",
    "            kernel_size=1,\n",
    "            stride=1,\n",
    "            padding=0\n",
    "        )\n",
    "\n",
    "    def make_image_noise(self, x):\n",
    "        batch_size, _, _, _ = x.size()\n",
    "        return torch.FloatTensor(batch_size, self.latent_dim, 1, 1).uniform_(0., 1.).to(x.device)\n",
    "    \n",
    "    def forward(self, input_noise):        \n",
    "        x = self.initial_block(input_noise)\n",
    "        x = self.initial_conv(x)\n",
    "\n",
    "        for block in self.blocks:\n",
    "            x = block(x)\n",
    "        \n",
    "        x = self.to_mri(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class DiscriminatorBlock(nn.Module):\n",
    "    def __init__(self, input_channels, filters, downsample=True):\n",
    "        super().__init__()\n",
    "        self.conv_res = nn.Conv2d(input_channels, filters, 1, stride = (2 if downsample else 1))\n",
    "\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Conv2d(input_channels, filters, 3, padding=1),\n",
    "            leaky_relu(),\n",
    "            nn.Conv2d(filters, filters, 3, padding=1),\n",
    "            leaky_relu()\n",
    "        )\n",
    "\n",
    "        self.downsample = nn.Sequential(\n",
    "            Blur(),\n",
    "            nn.Conv2d(filters, filters, 3, padding = 1, stride = 2)\n",
    "        ) if downsample else None\n",
    "\n",
    "    def forward(self, x):\n",
    "        res = self.conv_res(x)\n",
    "        x = self.net(x)\n",
    "        if exists(self.downsample):\n",
    "            x = self.downsample(x)\n",
    "        x = (x + res) * (1 / math.sqrt(2))\n",
    "        return x\n",
    "    \n",
    "    \n",
    "class Discriminator(nn.Module):\n",
    "    def __init__(self, \n",
    "                 image_size, \n",
    "                 network_capacity = 16, \n",
    "                 fmap_max = 512,\n",
    "                 cls=1):\n",
    "        super().__init__()\n",
    "        num_layers = int(log2(image_size) - 1)\n",
    "        num_init_filters = cls\n",
    "\n",
    "        blocks = []\n",
    "        filters = [num_init_filters] + [(network_capacity * 4) * (2 ** i) for i in range(num_layers + 1)]\n",
    "\n",
    "        set_fmap_max = partial(min, fmap_max)\n",
    "        filters = list(map(set_fmap_max, filters))\n",
    "        chan_in_out = list(zip(filters[:-1], filters[1:]))\n",
    "\n",
    "        blocks = []\n",
    "        quantize_blocks = []\n",
    "\n",
    "        for ind, (in_chan, out_chan) in enumerate(chan_in_out):\n",
    "            is_not_last = ind != (len(chan_in_out) - 1)\n",
    "\n",
    "            block = DiscriminatorBlock(in_chan, out_chan, downsample = is_not_last)\n",
    "            blocks.append(block)\n",
    "\n",
    "        self.blocks = nn.ModuleList(blocks)\n",
    "        self.quantize_blocks = nn.ModuleList(quantize_blocks)\n",
    "\n",
    "        chan_last = filters[-1]\n",
    "        latent_dim = 2 * 2 * chan_last\n",
    "\n",
    "        self.final_conv = nn.Conv2d(chan_last, chan_last, 3, padding=1)\n",
    "        self.flatten = Flatten()\n",
    "        self.to_logit = nn.Linear(latent_dim, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        b, *_ = x.shape\n",
    "\n",
    "        for block in self.blocks:\n",
    "            x = block(x)\n",
    "\n",
    "        x = self.final_conv(x)\n",
    "        x = self.flatten(x)\n",
    "        x = self.to_logit(x)\n",
    "        return x.squeeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen = Generator(256)\n",
    "x = torch.randn((2, 2, 256, 256))\n",
    "noise = gen.make_image_noise(x)\n",
    "print(noise.size())\n",
    "out = gen(noise)\n",
    "print(out.size())\n",
    "\n",
    "dis = Discriminator(256)\n",
    "out = dis(out)\n",
    "print(out)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
